# -*- coding: utf-8  -*-
# @Time : 2021/8/20  9:12
# @Author : zhangnengbo
# @File : split_trainAndtest.py
# @Company : HPY

import os
import random

# file_dir_mine = 'D:\\workspace\\data\\jindi\\mine'
# file_dir_waste = 'D:\\workspace\\data\\jindi\\waste'

file_dir_mine = 'D:\\data\\jindi\\SAVE\\mine'
file_dir_waste = 'D:\\data\\jindi\\SAVE\\waste'
files_mine = os.listdir(file_dir_mine)
files_waste = os.listdir(file_dir_waste)

sum_data = []
for i in files_mine:
    path = os.path.join(file_dir_mine, i)
    result = (path,1)
    sum_data.append(result)
    print(result)

for i in files_waste:
    path = os.path.join(file_dir_waste, i)
    result = (path, 0)
    sum_data.append(result)
    print(result)

random.shuffle(sum_data)

print('----------------------------------------------')

k = 0
for i in sum_data:
    name = i[0]
    label = i[1]
    content = name + ' ' + str(label) + '\n'
    if k%10 < 6:
        with open('D:\\data\\jindi\\SAVE\\JinDiTrain.txt', 'a') as f:
            f.write(content)
    else:
        with open('D:\\data\\jindi\\SAVE\\JinDiTest.txt', 'a') as f:
            f.write(content)
    k = k + 1



train = 1